// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================

/*!
\file
\brief Implementation of Stratified CMH Binary BT Association (BitStrings)
*/

#include <math.h>

#include "pcmh_assoc.h"
#include "pmaths.h"
#include "solfile.h"
#include "utils.h"

void pcmh_bt_assoc_write_file_header(FILE *out, struct selected_options *ops) {
	if (out != NULL && ops != NULL) {	
		int ci = (int) (100 * ops->ci_level);
		fprintf(out, "CHR,BP,SNP,CHISQ,P,OR,L%i,U%i,", ci, ci);
		if (ops->breslowday) fprintf(out, "CHISQ_BD,P_BD,");
		fprintf(out, "EMP1,EMP2,BONF,HOLM,FDR_BH\n");
	}
}

double pcmh_bt_assoc(double zt, int nLevels, BOOL X, BOOL haploid, BitString ones_and_twos, struct sample *samples, int nSamples, struct small_locus *snp, FILE *out, double emp1, double emp2, BOOL breslowday, struct adjusted *adj) {
	int i;
	double CMH = 0, denom = 0;

	// Disease X allele X strata

	// Calculate mean of 11 cell for each strata
	double mean_11[MAX_LEVELS];
	double var_11[MAX_LEVELS];

	// Calculate statistic
	double n_11[MAX_LEVELS];
	double n_12[MAX_LEVELS];
	double n_21[MAX_LEVELS];
	double n_22[MAX_LEVELS];

	// Disease marginals
	double n_1X[MAX_LEVELS]; // disease
	double n_2X[MAX_LEVELS]; // no disease

	double n_X1[MAX_LEVELS]; // F allele
	double n_X2[MAX_LEVELS]; // T allele
	double n_TT[MAX_LEVELS]; // Total allele count

	BOOL validK[MAX_LEVELS];

	// MH Odds ratio & CI
	double R = 0, S = 0, OR = 0;
	double r2[MAX_LEVELS], s2[MAX_LEVELS];
	double v1 = 0, v2 = 0, v3 = 0;
	double varOR, OR_lower, OR_upper;

	// Initialise the results arrays.
	zero_vector(mean_11, nLevels);
	zero_vector(var_11, nLevels);
	zero_vector(n_11, nLevels);
	zero_vector(n_12, nLevels);
	zero_vector(n_21, nLevels);
	zero_vector(n_22, nLevels);
	zero_vector(n_1X, nLevels);
	zero_vector(n_2X, nLevels);
	zero_vector(n_X1, nLevels);
	zero_vector(n_X2, nLevels);
	zero_vector(n_TT, nLevels);
	init_bool_array(validK, nLevels, FALSE);

	// Iterate over individuals
	for (i = 0; i < nSamples; i++) {
		unsigned int mask;
		long bitpos = (long) i * 2;
		BOOL s1 = FALSE, s2 = FALSE;
		struct sample *pperson =  samples[i].pperson;

		// Read S1 value
		mask = 0x1;
		mask = mask << (bitpos % BITBULKSIZE);
		mask = ones_and_twos[bitpos/BITBULKSIZE] & mask;
		if (mask !=0) s1 = TRUE;
		bitpos++;

		// Read S2 value
		mask = 0x1;
		mask = mask << (bitpos % BITBULKSIZE);
		mask = ones_and_twos[bitpos/BITBULKSIZE] & mask;
		if (mask !=0) s2 = TRUE;

		// Affected individuals
		if (pperson->aff && !pperson->missing) {
			// Haploid?
			if (haploid || (X && samples[i].sex)) {
				// Allelic marginal
				if (!s1) {
					// FF hom
					n_11[pperson->sol]++;
					n_X1[pperson->sol]++;		    
				} else  {
					if ( !s2) // FT
						continue;  // skip missing genotypes
					else {   // TT
						n_12[pperson->sol]++;
						n_X2[pperson->sol]++;		      
					}
				}

				// Disease marginal
				n_1X[pperson->sol]++;
				n_TT[pperson->sol]++;
			} else {  // autosomal
				// Allelic marginal
				if (!s1) {
					if (!s2) { // FF hom
						n_11[pperson->sol] += 2;
						n_X1[pperson->sol] += 2;
					} else {
						n_11[pperson->sol]++; // FT het
						n_12[pperson->sol]++;
						n_X1[pperson->sol]++;
						n_X2[pperson->sol]++;		      
					}
				} else {
					if (!s2) // FT
						continue;  // skip missing genotypes
					else { // TT
						n_12[ pperson->sol ] += 2;
						n_X2[ pperson->sol ] += 2;		      
					}
				}

				// Disease marginal
				n_1X[ pperson->sol ] += 2;
				n_TT[ pperson->sol ] += 2;
			} 
		}
		else if (!pperson->missing) { // Unaffecteds
			// Haploid?
			if (haploid || (X && samples[i].sex)) {
				// Allelic marginal
				if (!s1) {
					// FF hom
					n_21[pperson->sol]++;
					n_X1[pperson->sol]++;		    
				} else  {
					if (!s2) // FT
						continue;  // skip missing genotypes
					else { // TT
						n_22[pperson->sol]++;
						n_X2[pperson->sol]++;		      
					}
				}

				// Disease marginal
				n_2X[pperson->sol ]++;
				n_TT[pperson->sol ]++;

			} else {   // autosomal 
				// Allelic marginal
				if (!s1) {
					if (!s2) { // FF
						n_X1[pperson->sol] += 2;
						n_21[pperson->sol] += 2;
					} else {
						n_X1[pperson->sol]++;
						n_X2[pperson->sol]++;		      
						n_21[pperson->sol]++;
						n_22[pperson->sol]++;
					}
				} else  {
					if (!s2) // FT
						continue;  // skip missing genotypes
					else  { // TT
						n_X2[pperson->sol] += 2;		      		      
						n_22[pperson->sol] += 2;
					}
				}     

				// disease marginal
				n_2X[pperson->sol] += 2;
				n_TT[pperson->sol] += 2;
			}  
		}
	}

	// Finished iterating over individuals: cluster needs at least 2 
	// nonmissing individuals
	for (i = 0; i < nLevels; i++) if (n_TT[i] >= 2) validK[i] = TRUE;

	for (i = 0; i < nLevels; i++) {
		if (validK[i]) {
			mean_11[i] = (n_X1[i] * n_1X[i]) / n_TT[i] ;
			var_11[i] = (n_X1[i] * n_X2[i] * n_1X[i] * n_2X[i]) / (n_TT[i] * n_TT[i] * (n_TT[i]-1));
		}
	}

	for (i = 0; i < nLevels; i++) {
		if(validK[i]) {
			CMH += n_11[i] - mean_11[i];
			denom += var_11[i];
		}
	}

	CMH *= CMH;
	CMH /= denom;

	// MH Odds ratio & CI
	for (i = 0; i < nLevels; i++) {
		if (validK[i]) {
			r2[i] = (n_11[i] * n_22[i]) / n_TT[i];
			s2[i] = (n_12[i] * n_21[i]) / n_TT[i];
			R += r2[i]; 
			S += s2[i];
		}
	}
	OR = R / S;

	for (i = 0; i< nLevels; i++) {
		if (validK[i]) {
			v1 += (1/n_TT[i]) * (n_11[i] + n_22[i]) * r2[i];
			v2 += (1/n_TT[i]) * (n_12[i] + n_21[i]) * s2[i];
			v3 += (1/n_TT[i]) * ((n_11[i] + n_22[i]) * s2[i] + (n_12[i] + n_21[i]) * r2[i]);
		}
	}

	varOR = ((1/(2*R*R)) * v1) + ((1/(2*S*S)) * v2) + ((1/(2*R*S)) * v3);
	OR_lower = exp(log(OR) - (zt * sqrt(varOR)));
	OR_upper = exp(log(OR) + (zt * sqrt(varOR)));

	// Capture the CMH (a.k.a. CHISQ) statistic for the calculation to find FDR et cetera.
	if (adj != NULL && out == NULL) if (realnum(CMH)) adj->chisq = CMH;

	// Time to write stuff to file.
	if (snp != NULL && out != NULL && adj != NULL) {
		fprintf(out, "%i,%i,%s,", snp->chr, snp->bp, snp->name);

		if (realnum(CMH)) fprintf(out, "%10.16f,%10.16f,", CMH, chiprobP(CMH,1));
		else fprintf(out, "NA,NA,");

		if (realnum(OR)) fprintf(out, "%10.16f,", OR);
		else fprintf(out, "NA,");

		if (realnum(OR_lower)) fprintf(out, "%10.16f,", OR_lower);
		else fprintf(out, "NA,");

		if (realnum(OR_upper)) fprintf(out, "%10.16f,", OR_upper);
		else fprintf(out, "NA,");

		// Optional Breslow-Day test of homogeneity of odds ratios
		if (breslowday) {
			double amax, bb, determ, as_plus, as_minus, Astar, Bstar, Cstar, Dstar, Var, BDX2 = 0, BDp = 1.0;
			int df = 0;

			for (i = 0; i < nLevels; i++) {
				if (validK[i]) {
					df++;
					amax = (n_1X[i] < n_X1[i]) ? n_1X[i] : n_X1[i];
					bb = n_2X[i] + n_1X[i] * OR - n_X1[i] * (1-OR);
					determ = sqrt(bb*bb + 4*(1-OR) * OR * n_1X[i] * n_X1[i]);
					as_plus = ( -bb + determ ) / ( 2 - 2 * OR );
					as_minus = ( -bb - determ ) / ( 2 - 2 * OR );		      
					Astar =  as_minus <= amax && as_minus >= 0 ? as_minus  : as_plus ;
					Bstar = n_1X[i] - Astar;
					Cstar = n_X1[i] - Astar;
					Dstar = n_2X[i] - n_X1[i] + Astar;
					Var = 1/(1/Astar + 1/Bstar + 1/Cstar + 1/Dstar);
					BDX2 += ( (n_11[i] - Astar) * ( n_11[i] - Astar ) ) / Var ; 
				}
			}

			BDp = chiprobP(BDX2 , df-1); 

			if (BDp > -1) fprintf(out, "%10.16f,%10.16f,", BDX2, BDp);
			else fprintf(out, "NA,NA,");
		}

		fprintf(out, "%10.16f,%10.16f,%f,%f,%f\n", emp1, emp2, adj->bonf, adj->holm, adj->fdr_bh);
	}

	return CMH;
}